import os
import numpy as np
import torch as T
import torch.nn as nn
import torch.optim as optim
from torch.distributions.categorical import Categorical

class PPOMemory:
    def __init__(self,batch_size):
        self.states = [] # save the states as list form
        self.probs = []
        self.vals = [] # our critic calculates
        self.actions = []
        self.rewards = []
        self.dones = []

        self.batch_size = batch_size

    def generate_batches(self):
        n_states = len(self.states) # N = 20
        batch_start = np.arange(0,n_states,self.batch_size)
        indices = np.arange(n_states,dtype=np.int64)
        np.random.shuffle(indices)
        batches = [indices[i:i+self.batch_size] for i in batch_start]

        return np.array(self.states),\
                np.array(self.actions),\
                np.array(self.probs),\
                np.array(self.vals),\
                np.array(self.rewards),\
                np.array(self.dones),\
                batches

    def store_memory(self, state, action, probs, vals, reward, done):
        self.states.append(state)
        self.actions.append(action)
        self.probs.append(probs)
        self.vals.append(vals)
        self.rewards.append(reward)
        self.dones.append(done)

    def clear_memory(self):
        self.states = []
        self.probs = []
        self.actions = []
        self.rewards = []
        self.dones = []
        self.vals = []


class ActorNetwork(nn.Module):
    def __init__(self,n_actions,input_dims,alpha,
                 fc1_dims=64,fc2_dims=64,chkpt_dir='tmp/ppo'):
        super(ActorNetwork,self).__init__()

        self.checkpoint_file = os.path.join(chkpt_dir,'actor_torch_ppo')
        self.actor = nn.Sequential(
            nn.Linear(input_dims,fc1_dims), # 2
            nn.ReLU(),
            nn.Linear(fc1_dims,fc2_dims), # 256*256
            nn.ReLU(),
            nn.Linear(fc2_dims,n_actions),# 256*4, actions: up,down,right,left,
            nn.Softmax(dim=-1)
        )

        self.optimizer = optim.Adam(self.parameters(),lr=alpha)
        self.device = T.device('cpu')
        self.to(self.device)

    def forward(self,state):
        dist = self.actor(state)
        dist = Categorical(dist) # categorical表示创建probs为标准的类别样本分布，dist表示分布，dist.sample()表示抽取一次

        return dist

    def save_checkpoint(self):
        T.save(self.state_dict(),self.checkpoint_file)

    def load_checkpoint(self):
        self.load_state_dict(T.load(self.checkpoint_file))

class CriticNetwork(nn.Module):
    def __init__(self,input_dims,alpha,fc1_dims=64,fc2_dims=64,
                 chkpt_dir='tmp/ppo'):
        super(CriticNetwork,self).__init__()

        self.checkpoint_file = os.path.join(chkpt_dir,'critic_torch_ppo')
        self.critic = nn.Sequential(
            nn.Linear(input_dims,fc1_dims),
            nn.ReLU(),
            nn.Linear(fc1_dims,fc2_dims),
            nn.ReLU(),
            nn.Linear(fc2_dims,1)
        )

        self.optimizer = optim.Adam(self.parameters(),lr=alpha)
        self.device = T.device('cpu')
        self.to(self.device)

    def forward(self,state):
        value = self.critic(state)

        return value

    def save_checkpoint(self):
        T.save(self.state_dict(),self.checkpoint_file)

    def load_checkpoint(self):
        self.load_state_dict(T.load(self.checkpoint_file))




def del_tensor_ele(arr, index):
    arr1 = arr[0:index]
    arr2 = arr[index + 1:]
    return T.cat((arr1, arr2), dim=0)

class Agent:
    def __init__(self,n_actions,input_dims,gamma=0.99,alpha=0.001, gae_lamda=0.95,            # gamma, discount factor, alpha, learning rate, policy rate
                 policy_clip=0.2,batch_size=64,N=2048,n_epochs=10):
        self.gamma = gamma
        self.policy_clip = policy_clip
        self.n_epochs = n_epochs
        self.gae_lambda = gae_lamda

        self.actor = ActorNetwork(n_actions,input_dims,alpha)
        self.critic = CriticNetwork(input_dims,alpha)
        self.memory = PPOMemory(batch_size)

    def remember(self, state, action, probs, vals, reward, done):
        self.memory.store_memory(state, action, probs, vals, reward, done)

    def save_models(self):
        print('... saving models ...')
        self.actor.save_checkpoint()
        self.critic.save_checkpoint()

    def load_models(self):
        print('... loading models ...')
        self.actor.load_checkpoint()
        self.critic.load_checkpoint()

    def choose_action(self,observation):
        # Checking if the state exists in the table
        #self.check_state_exist(observation)
        #observation = list(map(float,observation))
        #epsilon = 0.7
        state = T.tensor(observation,dtype=T.float).to(self.actor.device)
        value = self.critic(state)
        dist = self.actor(state)  # give us our distribution for choosing an action
        # add a check list for invaild action
        #dist_re = dist.remove(2)
        logits = dist.logits
        probs_old = dist.probs
        logits_new = del_tensor_ele(logits,1)
        probs_new = del_tensor_ele(probs_old,1)

        dist_new = Categorical(probs_new)
        action2 = dist_new.sample().numpy()
        action = dist.sample() # sample our distribution
        # check the action is the invail action else input the sample data;

        probs = T.squeeze(dist.log_prob(action)).item()
        action = T.squeeze(action).item()
        value = T.squeeze(value).item()

        lkl=action



        return action,probs,value


    def learn(self):
        for _ in range(self.n_epochs):
            state_arr,action_arr,old_prob_arr,vals_arr,\
            reward_arr,dones_arr,batches = \
                self.memory.generate_batches()
            values = vals_arr
            advantage = np.zeros(len(reward_arr),dtype=np.float32)

            for t in range(len(reward_arr)-1):
                discount = 1
                a_t = 0
                for k in range(t,len(reward_arr)-1):
                    a_t += discount*(reward_arr[k]+self.gamma*values[k+1]* \
                                     (1-int(dones_arr[k]))-values[k])
                    discount *= self.gamma*self.gae_lambda
                advantage[t] = a_t
            advantage = T.tensor(advantage).to(self.actor.device)

            values = T.tensor(values).to(self.actor.device)
            for batch in batches:
                states = T.tensor(state_arr[batch], dtype=T.float).to(self.actor.device)
                old_probs = T.tensor(old_prob_arr[batch]).to(self.actor.device)
                actions = T.tensor(action_arr[batch]).to(self.actor.device)

                dist = self.actor(states)
                critic_value = self.critic(states)

                critic_value = T.squeeze(critic_value)

                new_probs = dist.log_prob(actions)
                prob_ratio = new_probs.exp() / old_probs.exp()
                # prob_ratio = (new_probs - old_probs).exp()
                weighted_probs = advantage[batch] * prob_ratio
                weighted_clipped_probs = T.clamp(prob_ratio, 1 - self.policy_clip,
                                                 1 + self.policy_clip) * advantage[batch]
                actor_loss = -T.min(weighted_probs, weighted_clipped_probs).mean()

                returns = advantage[batch] + values[batch]
                critic_loss = (returns - critic_value) ** 2
                critic_loss = critic_loss.mean()

                total_loss = actor_loss + 0.5 * critic_loss
                #print(f"------batch:",batch,"The total loss is:",total_loss.float())

                self.actor.optimizer.zero_grad()
                self.critic.optimizer.zero_grad()
                total_loss.backward()
                self.actor.optimizer.step()
                self.critic.optimizer.step()

        #self.memory.clear_memory()










